{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "30f483d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#计算某种配置下的内存开销和存储放大率\n",
    "pq_bucket_num = 32 #PQ的子空间数量, 例如32表示每个向量量化为32个Byte的向量\n",
    "raw_vector_dim = 128\n",
    "raw_vector_size = 1 #原始向量每一维的大小\n",
    "vector_num = 1000000\n",
    "\n",
    "L1_R = 32\n",
    "L1_nodes_ratio = 0.1\n",
    "L2_R = L1_R\n",
    "L3_Cluster = 32000\n",
    "redundancy = 2 #每个向量被分到多少个聚类中\n",
    "\n",
    "#以下是一些基本上不会有变化的变量\n",
    "pq_bit = 8\n",
    "pq_byte = pq_bit / 8\n",
    "id_size = 4 #每个节点id的字节数，一般用uint32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "74e18521",
   "metadata": {},
   "outputs": [],
   "source": [
    "#计算内存开销\n",
    "L1_mem = L1_nodes_ratio * L3_Cluster * (pq_bucket_num + L1_R*id_size) #L1需要常驻点的邻居及其PQ向量\n",
    "L2_mem = raw_vector_dim * raw_vector_size * 2**pq_bit #L2仅PQ表需要常驻内存\n",
    "L3_mem = L3_Cluster * id_size #L3仅偏移列表需要常驻内存\n",
    "Total_mem = L1_mem + L2_mem + L3_mem\n",
    "\n",
    "#计算存储放大率.这里都不考虑块对齐的问题\n",
    "L1_space = L1_mem #L1的空间开销就是内存索引的存储开销\n",
    "L2_space = L2_mem + L3_Cluster*L2_R*(id_size+pq_bucket_num)#L2的空间开销就是图的大小\n",
    "L3_space = L3_mem + redundancy*vector_num*raw_vector_size*raw_vector_dim\n",
    "Total_space = L1_space + L2_space + L3_space\n",
    "ratio = Total_space / (vector_num*raw_vector_size*raw_vector_dim)\n",
    "#计算nodetype4下的存储放大率\n",
    "L2_space_wraw = L2_mem + L3_Cluster*(L2_R*(id_size+pq_bucket_num)+raw_vector_dim*raw_vector_size)#每个节点额外存储原始向量\n",
    "Total_space_wraw = L1_space + L2_space_wraw + L3_space\n",
    "ratio_wraw = Total_space_wraw / (vector_num*raw_vector_size*raw_vector_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "7815e688",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L1_mem: 422400.0 B, 412.5 KB, 0.40283203125 MB\n",
      "L2_mem: 32768 B, 32.0 KB, 0.03125 MB\n",
      "L3_mem: 128000 B, 125.0 KB, 0.1220703125 MB\n",
      "Total mem: 583168.0 B, 569.5 KB, 0.55615234375 MB\n",
      "\n",
      "L1_space: 422400.0 B, 412.5 KB, 0.40283203125 MB\n",
      "L2_space: 8224768 B, 8032.0 KB, 7.84375 MB\n",
      "L3_space: 256128000 B, 250125.0 KB, 244.2626953125 MB\n",
      "Total space: 264775168.0 B, 258569.5 KB, 252.50927734375 MB\n",
      "ratio: 2.068556\n",
      "\n",
      "If include raw vector in L2:\n",
      "L2_space_wraw: 12320768 B, 12032.0 KB, 11.75 MB\n",
      "Total_space_wraw: 268871168.0 B, 262569.5 KB, 256.41552734375 MB\n",
      "ratio_wraw: 2.100556\n"
     ]
    }
   ],
   "source": [
    "print(\"L1_mem:\",L1_mem,\"B,\", L1_mem/1024,\"KB,\", L1_mem/1024/1024,\"MB\")\n",
    "print(\"L2_mem:\",L2_mem,\"B,\", L2_mem/1024,\"KB,\", L2_mem/1024/1024,\"MB\")\n",
    "print(\"L3_mem:\",L3_mem,\"B,\", L3_mem/1024,\"KB,\", L3_mem/1024/1024,\"MB\")\n",
    "print(\"Total mem:\",Total_mem,\"B,\", Total_mem/1024,\"KB,\", Total_mem/1024/1024,\"MB\")\n",
    "print(\"\")\n",
    "print(\"L1_space:\",L1_space,\"B,\", L1_space/1024,\"KB,\", L1_space/1024/1024,\"MB\")\n",
    "print(\"L2_space:\",L2_space,\"B,\", L2_space/1024,\"KB,\", L2_space/1024/1024,\"MB\")\n",
    "print(\"L3_space:\",L3_space,\"B,\", L3_space/1024,\"KB,\", L3_space/1024/1024,\"MB\")\n",
    "print(\"Total space:\",Total_space,\"B,\", Total_space/1024,\"KB,\", Total_space/1024/1024,\"MB\")\n",
    "print(\"ratio:\",ratio)\n",
    "print(\"\\nIf include raw vector in L2:\")\n",
    "print(\"L2_space_wraw:\",L2_space_wraw,\"B,\", L2_space_wraw/1024,\"KB,\", L2_space_wraw/1024/1024,\"MB\")\n",
    "print(\"Total_space_wraw:\",Total_space_wraw,\"B,\", Total_space_wraw/1024,\"KB,\", Total_space_wraw/1024/1024,\"MB\")\n",
    "print(\"ratio_wraw:\",ratio_wraw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "234caa2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DiskANN_mem: 4032768 B, 3938.25 KB, 3.845947265625 MB\n",
      "DiskANN_space: 256000000 B, 250000.0 KB, 244.140625 MB\n",
      "ratio: 2.0\n"
     ]
    }
   ],
   "source": [
    "#计算DiskANN的开销\n",
    "DiskANN_mem = L2_mem + pq_bucket_num*vector_num #DiskANN主要就是内存存储PQ向量及其码表的开销\n",
    "DiskANN_space = vector_num*(L2_R*id_size+raw_vector_size*raw_vector_dim)#DiskANN的存储开销主要就是存储原始向量以及邻居的开销\n",
    "print(\"DiskANN_mem:\",DiskANN_mem,\"B,\", DiskANN_mem/1024,\"KB,\", DiskANN_mem/1024/1024,\"MB\")\n",
    "print(\"DiskANN_space:\",DiskANN_space,\"B,\", DiskANN_space/1024,\"KB,\", DiskANN_space/1024/1024,\"MB\")\n",
    "print(\"ratio:\",DiskANN_space/(vector_num*raw_vector_size*raw_vector_dim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "54f77fec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SPANN_mem: 4256768 B, 4157.0 KB, 4.0595703125 MB\n",
      "SPANN_space: 260256768 B, 254157.0 KB, 248.2001953125 MB\n",
      "ratio: 2.033256\n"
     ]
    }
   ],
   "source": [
    "#计算SPANN的开销，这里假设SPANN也用PQ向量\n",
    "SPANN_mem = L2_mem + L3_Cluster * (pq_bucket_num + L1_R*id_size) #SPANN主要就是内存存图的开销和PQ表的开销\n",
    "SPANN_space = SPANN_mem + redundancy*vector_num*raw_vector_size*raw_vector_dim #SPANN的存储开销就是冗余存储的开销\n",
    "print(\"SPANN_mem:\",SPANN_mem,\"B,\", SPANN_mem/1024,\"KB,\", SPANN_mem/1024/1024,\"MB\")\n",
    "print(\"SPANN_space:\",SPANN_space,\"B,\", SPANN_space/1024,\"KB,\", SPANN_space/1024/1024,\"MB\")\n",
    "print(\"ratio:\",SPANN_space/(vector_num*raw_vector_size*raw_vector_dim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "c3486504",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AiSAQ_mem: 32768 B, 32.0 KB, 0.03125 MB\n",
      "AiSAQ_space_without_raw: 256032768 B, 250032.0 KB, 244.171875 MB\n",
      "ratio_without_raw: 2.000256\n",
      "AiSAQ_space_with_raw: 384032768 B, 375032.0 KB, 366.2421875 MB\n",
      "ratio_with_raw: 3.000256\n"
     ]
    }
   ],
   "source": [
    "#计算AiSAQ的开销\n",
    "AiSAQ_mem = L2_mem #AiSAQ只需要常驻PQ表即可\n",
    "AiSAQ_space_woraw = AiSAQ_mem + vector_num*L2_R*(id_size+pq_bucket_num)  #不包含原始向量的版本，精度低存储消耗少\n",
    "AiSAQ_space_wraw = AiSAQ_mem + vector_num*L2_R*(id_size+pq_bucket_num) + vector_num*raw_vector_size*raw_vector_dim#包含原始向量的版本，精度高存储消耗大\n",
    "print(\"AiSAQ_mem:\",AiSAQ_mem,\"B,\", AiSAQ_mem/1024,\"KB,\", AiSAQ_mem/1024/1024,\"MB\")\n",
    "print(\"AiSAQ_space_without_raw:\",AiSAQ_space_woraw,\"B,\", AiSAQ_space_woraw/1024,\"KB,\", AiSAQ_space_woraw/1024/1024,\"MB\")\n",
    "print(\"ratio_without_raw:\",AiSAQ_space_woraw/(vector_num*raw_vector_size*raw_vector_dim))\n",
    "print(\"AiSAQ_space_with_raw:\",AiSAQ_space_wraw,\"B,\", AiSAQ_space_wraw/1024,\"KB,\", AiSAQ_space_wraw/1024/1024,\"MB\")\n",
    "print(\"ratio_with_raw:\",AiSAQ_space_wraw/(vector_num*raw_vector_size*raw_vector_dim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9a28176",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
